#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
# @File   : trt_infer.py
# @Author : yuanwenjin
# @Mail   : xxxx@mail.com
# @Date   : 2020/06/30 18:46:49
# @Docs   : 使用tensorRT进行推理
'''

import numpy as np
import tensorrt as trt
from trt_utils import load_engine, TRT_LOGGER
import pycuda.driver as cuda
import pycuda.autoinit
import torch

# Simple helper data class that's a little nicer to use than a 2-tuple.
class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        self.host = host_mem
        self.device = device_mem

    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        return self.__str__()

# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
def allocate_buffers(engine, batch_size=1):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()
    for binding in engine:
        # size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        size = trt.volume(engine.get_binding_shape(binding)) * batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # Allocate host and device buffers
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        # Append the device buffer to device bindings.
        bindings.append(int(device_mem))
        # Append to the appropriate list.
        if engine.binding_is_input(binding):
            inputs.append(HostDeviceMem(host_mem, device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem, device_mem))
    return inputs, outputs, bindings, stream

def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
    # Transfer input data to the GPU.
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    # Run inference.
    context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
    # Transfer predictions back from the GPU.
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # Synchronize the stream
    stream.synchronize()
    # Return only the host outputs.
    return [out.host for out in outputs]

class trt_infer(object):
    '''使用tensorRT进行推理的类, 要求输入图像大小都相同, 且为生成engine时设置大小, 否则有问题'''
    def __init__(self, engine_path=None, batch_size=1):
        '''初始化engine'''
        if engine_path is None:
            raise ValueError('engine is None!')
        self.batch_size = batch_size
        self.engine_path = engine_path
        self.engine = load_engine(engine_path)
        self.context = self.engine.create_execution_context()
        self.buffers = allocate_buffers(self.engine, batch_size)

    def infer(self, ims):
        '''进行推理'''
        im_len = len(ims)
        if im_len != self.batch_size:
            # 需根据输入图像数量分配空间
            self.batch_size = im_len
            self.buffers = allocate_buffers(self.engine, im_len)
        inputs, outputs, bindings, stream = self.buffers
        np.copyto(inputs[0].host, np.array(ims).ravel())
        trt_outputs = do_inference(context=self.context,
                                    bindings=bindings,
                                    inputs=inputs,
                                    outputs=outputs,
                                    stream=stream,
                                    batch_size=self.batch_size)
        return trt_outputs

class trt_infer_object(object):
    '''使用tensorRT进行推理的类, 用于目标识别, 图像大小可以不相同, 但不能超过max_batch_size空间'''
    def __init__(self, engine_path=None, batch_size=1):
        '''初始化engine'''
        if engine_path is None:
            raise ValueError('engine is None!')
        self.batch_size = batch_size
        self.engine_path = engine_path
        self.engine = load_engine(engine_path)

    def infer(self, ims, class_num):
        '''进行推理'''
        im_len = len(ims)
        inputs, outputs, bindings, stream = allocate_buffers(self.engine, im_len)
        with self.engine.create_execution_context() as context:
            img = torch.from_numpy(np.array(ims)).float().numpy()
            inputs[0].host = img
            [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
            stream.synchronize()
            context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
            stream.synchronize()
            [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
            stream.synchronize()

            return [np.array(outputs[0].host.reshape(im_len, -1, class_num))]
            # return [out.host for out in outputs]
